{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install tensorflow_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "# @Time : 2020/4/22 21:46\n",
    "# @Author : zdqzyx\n",
    "# @File : main.py\n",
    "# @Software: PyCharm\n",
    "import sys\n",
    "sys.path.append('./')\n",
    "\n",
    "import tensorflow_datasets as tfds\n",
    "import tensorflow as tf\n",
    "import time\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from modeling import Transformer\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "examples, metadata = tfds.load('ted_hrlr_translate/pt_to_en', with_info=True,\n",
    "                               as_supervised=True)\n",
    "train_examples, val_examples = examples['train'], examples['validation']\n",
    "\n",
    "tokenizer_en = tfds.features.text.SubwordTextEncoder.build_from_corpus(\n",
    "    (en.numpy() for pt, en in train_examples), target_vocab_size=2**13)\n",
    "\n",
    "tokenizer_pt = tfds.features.text.SubwordTextEncoder.build_from_corpus(\n",
    "    (pt.numpy() for pt, en in train_examples), target_vocab_size=2**13)\n",
    "\n",
    "\n",
    "def encode(lang1, lang2):\n",
    "    lang1 = [tokenizer_pt.vocab_size] + tokenizer_pt.encode(\n",
    "        lang1.numpy()) + [tokenizer_pt.vocab_size + 1]\n",
    "\n",
    "    lang2 = [tokenizer_en.vocab_size] + tokenizer_en.encode(\n",
    "        lang2.numpy()) + [tokenizer_en.vocab_size + 1]\n",
    "\n",
    "    return lang1, lang2\n",
    "\n",
    "MAX_LENGTH = 40\n",
    "BUFFER_SIZE = 20000\n",
    "BATCH_SIZE = 64\n",
    "\n",
    "def filter_max_length(x, y, max_length=MAX_LENGTH):\n",
    "    return tf.logical_and(tf.size(x) <= max_length,\n",
    "                        tf.size(y) <= max_length)\n",
    "def tf_encode(pt, en):\n",
    "    result_pt, result_en = tf.py_function(encode, [pt, en], [tf.int64, tf.int64])\n",
    "    result_pt.set_shape([None])\n",
    "    result_en.set_shape([None])\n",
    "    return result_pt, result_en\n",
    "\n",
    "train_preprocessed = (\n",
    "    train_examples\n",
    "    .map(tf_encode)\n",
    "    .filter(filter_max_length)\n",
    "    # cache the dataset to memory to get a speedup while reading from it.\n",
    "    .cache()\n",
    "    .shuffle(BUFFER_SIZE))\n",
    "\n",
    "val_preprocessed = (\n",
    "    val_examples\n",
    "    .map(tf_encode)\n",
    "    .filter(filter_max_length))\n",
    "\n",
    "train_dataset = (train_preprocessed\n",
    "                 .padded_batch(BATCH_SIZE, padded_shapes=([None], [None]))\n",
    "                 .prefetch(tf.data.experimental.AUTOTUNE))\n",
    "\n",
    "\n",
    "val_dataset = (val_preprocessed\n",
    "               .padded_batch(BATCH_SIZE,  padded_shapes=([None], [None])))\n",
    "\n",
    "\n",
    "\n",
    "def checkout_dir(dir_path, do_delete=False):\n",
    "    import shutil\n",
    "    if do_delete and os.path.exists(dir_path):\n",
    "        shutil.rmtree(dir_path)\n",
    "    if not os.path.exists(dir_path):\n",
    "        print(dir_path, 'make dir ok')\n",
    "        os.makedirs(dir_path)\n",
    "\n",
    "class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
    "    def __init__(self, d_model, warmup_steps=4000):\n",
    "        super(CustomSchedule, self).__init__()\n",
    "\n",
    "        self.d_model = d_model\n",
    "        self.d_model = tf.cast(self.d_model, tf.float32)\n",
    "\n",
    "        self.warmup_steps = warmup_steps\n",
    "\n",
    "    def __call__(self, step):\n",
    "        arg1 = tf.math.rsqrt(step)\n",
    "        arg2 = step * (self.warmup_steps ** -1.5)\n",
    "\n",
    "        return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)\n",
    "\n",
    "\n",
    "def create_padding_mask(seq):\n",
    "    seq = tf.cast(tf.math.equal(seq, 0), tf.float32)\n",
    "    # 添加额外的维度来将填充加到\n",
    "    # 注意力对数（logits）。\n",
    "    return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)\n",
    "\n",
    "def create_look_ahead_mask(size):\n",
    "    '''\n",
    "    eg.\n",
    "    x = tf.random.uniform((1, 3))\n",
    "    temp = create_look_ahead_mask(x.shape[1])\n",
    "    temp:<tf.Tensor: shape=(3, 3), dtype=float32, numpy=\n",
    "            array([[0., 1., 1.],\n",
    "                   [0., 0., 1.],\n",
    "                   [0., 0., 0.]], dtype=float32)>\n",
    "    '''\n",
    "    mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n",
    "    return mask  # (seq_len, seq_len)\n",
    "\n",
    "def create_masks(inp, tar):\n",
    "    # 编码器填充遮挡\n",
    "    enc_padding_mask = create_padding_mask(inp)\n",
    "    # 在解码器的第二个注意力模块使用。\n",
    "    # 该填充遮挡用于遮挡编码器的输出。\n",
    "    dec_padding_mask = create_padding_mask(inp)\n",
    "    # 在解码器的第一个注意力模块使用。\n",
    "    # 用于填充（pad）和遮挡（mask）解码器获取到的输入的后续标记（future tokens）。\n",
    "    look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1]) #(tar_seq_len, tar_seq_len)\n",
    "    dec_target_padding_mask = create_padding_mask(tar) # (batch_size, 1, 1, tar_seq_len)\n",
    "    # 广播机制，look_ahead_mask==>(batch_size, 1, tar_seq_len, tar_seq_len)\n",
    "    # dec_target_padding_mask ==> (batch_size, 1, tar_seq_len, tar_seq_len)\n",
    "    combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n",
    "    return enc_padding_mask, combined_mask, dec_padding_mask\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_vocab_size is 8216, target_vocab_size is 8089\n"
     ]
    }
   ],
   "source": [
    "# ==============================================================\n",
    "\n",
    "params = {\n",
    "    'num_layers':4,\n",
    "    'd_model':128,\n",
    "    'dff':512,\n",
    "    'num_heads':8,\n",
    "    'input_vocab_size' :tokenizer_pt.vocab_size + 2,\n",
    "    'target_vocab_size':tokenizer_en.vocab_size + 2,\n",
    "    'pe_input':tokenizer_pt.vocab_size + 2,\n",
    "    'pe_target':tokenizer_en.vocab_size + 2,\n",
    "    'rate':0.1,\n",
    "}\n",
    "\n",
    "\n",
    "print('input_vocab_size is {}, target_vocab_size is {}'.format(params['input_vocab_size'], params['target_vocab_size']))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Latest checkpoint restored!!\n",
      "tf.Tensor([8087   16   13    7  328   10   14   24    5  966   19    2], shape=(12,), dtype=int32)\n",
      "Input: este é um problema que temos que resolver.\n",
      "Predicted translation: this is a problem that we have to solve it .\n",
      "Real translation: this is a problem we have to solve .\n",
      "tf.Tensor([8087   32 4220 7863  505   35   16  253   35   16  508    2], shape=(12,), dtype=int32)\n",
      "Input: os meus vizinhos ouviram sobre esta ideia.\n",
      "Predicted translation: my neighbors heard about this idea about this two .\n",
      "Real translation: and my neighboring homes heard about this idea .\n",
      "tf.Tensor([8087   16   13    3  124  774   10   12   98   19    2], shape=(11,), dtype=int32)\n",
      "Input: este é o primeiro livro que eu fiz.\n",
      "Predicted translation: this is the first book that i did it .\n",
      "Real translation: this is the first book i've ever done.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "class ModelHelper:\n",
    "\n",
    "    def __init__(self):\n",
    "        self.transformer  = Transformer(params)\n",
    "        # optimizer\n",
    "        learning_rate = CustomSchedule(params['d_model'])\n",
    "        self.optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)\n",
    "        self.loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')\n",
    "\n",
    "        # 主要为了累计一个epoch中的batch的loss，最后求平均，得到一个epoch的loss\n",
    "        self.train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
    "        # 主要为了累计一个epoch中的batch的acc，最后求平均，得到一个epoch的acc\n",
    "        self.train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')\n",
    "\n",
    "        self.test_loss = tf.keras.metrics.Mean(name='test_loss')\n",
    "        self.test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')\n",
    "\n",
    "        # 检查点\n",
    "        checkpoint_path = \"./checkpoints/train\"\n",
    "#         checkout_dir(dir_path=checkpoint_path, do_delete=True)\n",
    "        ckpt = tf.train.Checkpoint(transformer=self.transformer,\n",
    "                                   optimizer=self.optimizer)\n",
    "        ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)\n",
    "        # 如果检查点存在，则恢复最新的检查点。\n",
    "        if ckpt_manager.latest_checkpoint:\n",
    "            ckpt.restore(ckpt_manager.latest_checkpoint)\n",
    "            print('Latest checkpoint restored!!')\n",
    "\n",
    "    def loss_function(self, real, pred):\n",
    "        mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
    "        loss_ = self.loss_object(real, pred)\n",
    "        mask = tf.cast(mask, dtype=loss_.dtype)\n",
    "        loss_ *= mask\n",
    "        return tf.reduce_mean(loss_)\n",
    "\n",
    "\n",
    "    train_step_signature = [\n",
    "        tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
    "        tf.TensorSpec(shape=(None, None), dtype=tf.int64),\n",
    "    ]\n",
    "    @tf.function(input_signature=train_step_signature)\n",
    "    def train_step(self, inp, tar):\n",
    "        tar_inp = tar[:, :-1]\n",
    "        tar_real = tar[:, 1:]\n",
    "\n",
    "        enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions, _ = self.transformer(inp, tar_inp,\n",
    "                                         True,\n",
    "                                         enc_padding_mask,\n",
    "                                         combined_mask,\n",
    "                                         dec_padding_mask)\n",
    "            loss = self.loss_function(tar_real, predictions)\n",
    "\n",
    "        gradients = tape.gradient(loss, self.transformer.trainable_variables)\n",
    "        self.optimizer.apply_gradients(zip(gradients, self.transformer.trainable_variables))\n",
    "        self.train_loss(loss)\n",
    "        self.train_accuracy(tar_real, predictions)\n",
    "\n",
    "    @tf.function\n",
    "    def test_step(self, inp, labels):\n",
    "        predictions = self.predict(inp)\n",
    "        t_loss = self.loss_object(labels, predictions)\n",
    "        self.test_loss(t_loss)\n",
    "        self.test_accuracy(labels, predictions)\n",
    "\n",
    "    def train(self, train_dataset):\n",
    "        for epoch in range(params['epochs']):\n",
    "            start = time.time()\n",
    "            self.train_loss.reset_states()\n",
    "            self.train_accuracy.reset_states()\n",
    "            # inp -> portuguese, tar -> english\n",
    "            for (batch, (inp, tar)) in enumerate(train_dataset):\n",
    "                self.train_step(inp, tar)\n",
    "                if batch % 50 == 0:\n",
    "                    print('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1, batch, self.train_loss.result(), self.train_accuracy.result()))\n",
    "            if (epoch + 1) % 5 == 0:\n",
    "                ckpt_save_path = self.ckpt_manager.save()\n",
    "                print('Saving checkpoint for epoch {} at {}'.format(epoch + 1,ckpt_save_path))\n",
    "            print('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1, self.train_loss.result(), self.train_accuracy.result()))\n",
    "            print('Time taken for 1 epoch: {} secs\\n'.format(time.time() - start))\n",
    "\n",
    "    # 评估\n",
    "    def predict(self, inp_sentence):\n",
    "        start_token = [tokenizer_pt.vocab_size]\n",
    "        end_token = [tokenizer_pt.vocab_size + 1]\n",
    "\n",
    "        # 输入语句是葡萄牙语，增加开始和结束标记\n",
    "        inp_sentence = start_token + tokenizer_pt.encode(inp_sentence) + end_token\n",
    "        encoder_input = tf.expand_dims(inp_sentence, 0)\n",
    "\n",
    "        # 因为目标是英语，输入 transformer 的第一个词应该是\n",
    "        # 英语的开始标记。\n",
    "        decoder_input = [tokenizer_en.vocab_size]\n",
    "        output = tf.expand_dims(decoder_input, 0)\n",
    "\n",
    "        for i in range(MAX_LENGTH):\n",
    "            enc_padding_mask, combined_mask, dec_padding_mask = create_masks(\n",
    "                encoder_input, output)\n",
    "\n",
    "            # predictions.shape == (batch_size, seq_len, vocab_size)\n",
    "            predictions, attention_weights = self.transformer(encoder_input,\n",
    "                                                         output,\n",
    "                                                         False,\n",
    "                                                         enc_padding_mask,\n",
    "                                                         combined_mask,\n",
    "                                                         dec_padding_mask)\n",
    "\n",
    "            # 从 seq_len 维度选择最后一个词\n",
    "            predictions = predictions[:, -1:, :]  # (batch_size, 1, vocab_size)\n",
    "            predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)\n",
    "            # 如果 predicted_id 等于结束标记，就返回结果\n",
    "            if predicted_id == tokenizer_en.vocab_size + 1:\n",
    "                return tf.squeeze(output, axis=0), attention_weights\n",
    "            # 连接 predicted_id 与输出，作为解码器的输入传递到解码器。\n",
    "            output = tf.concat([output, predicted_id], axis=-1)\n",
    "        return tf.squeeze(output, axis=0)\n",
    "\n",
    "\n",
    "    def evaluation(self, x, y):\n",
    "        self.test_accuracy.reset_states()\n",
    "        self.test_loss.reset_states()\n",
    "        self.test_step(x, y, training=False)\n",
    "        template = 'Evaluation || Loss: {}, Accuracy: {}'\n",
    "        print(template.format(self.test_loss.result(), self.test_accuracy.result() * 100))\n",
    "\n",
    "\n",
    "def translate(model_helper, sentence):\n",
    "    result, _ = model_helper.predict(sentence)\n",
    "    print(result)\n",
    "    predicted_sentence = tokenizer_en.decode([i for i in result if i < tokenizer_en.vocab_size])\n",
    "    print('Input: {}'.format(sentence))\n",
    "    print('Predicted translation: {}'.format(predicted_sentence))\n",
    "\n",
    "model_helper = ModelHelper()\n",
    "\n",
    "\n",
    "translate(model_helper, \"este é um problema que temos que resolver.\")\n",
    "print (\"Real translation: this is a problem we have to solve .\")\n",
    "\n",
    "\n",
    "translate(model_helper, \"os meus vizinhos ouviram sobre esta ideia.\")\n",
    "print (\"Real translation: and my neighboring homes heard about this idea .\")\n",
    "\n",
    "\n",
    "translate(model_helper, \"este é o primeiro livro que eu fiz.\",)\n",
    "print (\"Real translation: this is the first book i've ever done.\")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:py37_tf20_gpu]",
   "language": "python",
   "name": "conda-env-py37_tf20_gpu-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
